import json
import idc
import idaapi
import idautils

def apply_il2cpp_data():
    """
    解析Il2CppDumper v6.7.46生成的script.json文件并应用到IDA Pro
    """
    
    # 读取script.json文件 - 修改为你的实际路径
    json_file_path = r'C:\Users\21585\Desktop\compare\script.json'
    
    try:
        with open(json_file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
    except Exception as e:
        print(f"Error reading JSON file: {e}")
        return
    
    # 计算基地址偏移
    image_base = idaapi.get_imagebase()
    print(f"IDA Image Base: 0x{image_base:X}")
    
    # 获取段信息
    segments = []
    for seg in idautils.Segments():
        seg_name = idc.get_segm_name(seg)
        seg_start = idc.get_segm_start(seg)
        seg_end = idc.get_segm_end(seg)
        segments.append((seg_name, seg_start, seg_end))
    
    if segments:
        print("Available segments:")
        for name, start, end in segments:
            print(f"  {name}: 0x{start:X} - 0x{end:X}")
    
    # 尝试找到.text段或主要代码段
    text_base = None
    for name, start, end in segments:
        if '.text' in name.lower() or 'code' in name.lower() or name == '':
            text_base = start
            break
    
    if text_base is None and segments:
        text_base = segments[0][1]  # 使用第一个段的基地址
    
    print(f"Using text segment base: 0x{text_base:X}")
    
    # 计算偏移量
    base_offset = text_base
    
    # 应用方法名和签名
    methods = data.get('ScriptMethod', [])
    print(f"Found {len(methods)} methods")
    
    # 用于跟踪重复名称
    name_counters = {}
    success_count = 0
    
    for i, method in enumerate(methods):
        try:
            addr = method.get('Address')
            name = method.get('Name')
            signature = method.get('Signature', '')
            type_sig = method.get('TypeSignature', '')
            
            if addr and name:
                # 转换地址：将RVA转换为IDA中的实际地址
                ida_addr = addr + base_offset
                
                # 检查地址是否在有效范围内
                if not idaapi.is_loaded(ida_addr):
                    continue
                
                # 清理函数名，替换IDA不支持的字符
                base_clean_name = name.replace('$', '_').replace('<', '_').replace('>', '_').replace('|', '_').replace('::', '_').replace('.', '_')
                
                # 处理重复名称
                if base_clean_name in name_counters:
                    name_counters[base_clean_name] += 1
                    clean_name = f"{base_clean_name}_{name_counters[base_clean_name]}"
                else:
                    name_counters[base_clean_name] = 0
                    clean_name = base_clean_name
                
                # 确保名称不超过IDA的限制（通常是512字符）
                if len(clean_name) > 200:
                    clean_name = clean_name[:200] + f"_{addr:X}"
                
                # 设置函数名
                result = idc.set_name(ida_addr, clean_name)
                if result:
                    success_count += 1
                    
                    # 设置函数注释包含完整签名
                    if signature:
                        idc.set_func_cmt(ida_addr, signature, 0)
                    
                    # 可选：设置重复注释包含类型签名
                    if type_sig:
                        idc.set_func_cmt(ida_addr, f"{signature}\nType: {type_sig}", 1)
                else:
                    # 如果还是失败，尝试使用地址作为后缀
                    fallback_name = f"{base_clean_name}_{addr:X}"
                    if idc.set_name(ida_addr, fallback_name):
                        success_count += 1
                        print(f"Used fallback name for 0x{ida_addr:X}: {fallback_name}")
                    else:
                        print(f"Failed to set name for address 0x{ida_addr:X}: {clean_name}")
                
        except Exception as e:
            print(f"Error processing method {i}: {e}")
            continue
        
        # 显示进度
        if (i + 1) % 1000 == 0:
            print(f"Processed {i + 1}/{len(methods)} methods...")
    
    print(f"Successfully applied {success_count}/{len(methods)} method names")
    print(f"Found {len(name_counters)} unique base names with {sum(v for v in name_counters.values() if v > 0)} duplicates")
    
    # 应用字符串数据
    if 'ScriptString' in data:
        strings = data.get('ScriptString', [])
        print(f"Found {len(strings)} strings")
        
        string_success = 0
        for i, string_data in enumerate(strings):
            try:
                # 字符串可能是对象格式或数组格式，需要适配
                if isinstance(string_data, dict):
                    addr = string_data.get('Address')
                    string_val = string_data.get('Value', '')
                elif isinstance(string_data, list) and len(string_data) >= 2:
                    addr = string_data[0]
                    string_val = string_data[1]
                else:
                    continue
                
                if addr and string_val:
                    # 转换地址
                    ida_addr = addr + base_offset
                    
                    # 检查地址是否在有效范围内
                    if not idaapi.is_loaded(ida_addr):
                        continue
                    
                    # 创建字符串
                    str_len = len(string_val.encode('utf-8'))
                    if idc.create_strlit(ida_addr, ida_addr + str_len + 1):
                        string_success += 1
                        
            except Exception as e:
                print(f"Error processing string {i}: {e}")
                continue
        
        print(f"Successfully applied {string_success}/{len(strings)} strings")
    
    # 应用其他数据类型
    if 'ScriptMetadata' in data:
        metadata = data.get('ScriptMetadata', [])
        print(f"Found {len(metadata)} metadata entries")
        
        for meta in metadata:
            try:
                addr = meta.get('Address')
                name = meta.get('Name', '')
                
                if addr and name:
                    # 转换地址
                    ida_addr = addr + base_offset
                    
                    if idaapi.is_loaded(ida_addr):
                        clean_name = name.replace('$', '_').replace('<', '_').replace('>', '_')
                        idc.set_name(ida_addr, f"metadata_{clean_name}")
                        
            except Exception as e:
                print(f"Error processing metadata: {e}")
                continue

def main():
    """
    主函数
    """
    print("Starting Il2Cpp data import...")
    print("=" * 50)
    
    # 确保IDA已经完成初始分析
    idaapi.auto_wait()
    
    # 应用Il2Cpp数据
    apply_il2cpp_data()
    
    print("=" * 50)
    print("Il2Cpp data import completed!")
    print("You may need to refresh IDA views (View -> Refresh)")

# 运行脚本
if __name__ == "__main__":
    main()